import numpy as np

from DeepSORT.algorithm.linear_assignment import *
from utils.IoU_utils import get_IoU


# 注意，IoU 操作的bbox的输入要求是xywh，但是网络输出的格式是xyxy

def IoU(bbox, candidates):
    IoU_list = []
    for candidate in candidates:
        IoU_list.append(get_IoU(bbox, candidate))
    return np.array(IoU_list)


def IoU_cost(tracks, detections,
             track_indices=None, detection_indices=None):
    """
    Returns
    -------
    ndarray
        Returns a cost matrix of shape
        len(track_indices), len(detection_indices) where entry (i, j) is
        `1 - iou(tracks[track_indices[i]], detections[detection_indices[j]])`.
    """
    if track_indices is None:
        track_indices = np.arange(len(tracks))
    if detection_indices is None:
        detection_indices = np.arange(len(detections))

    cost_matrix = np.zeros(
        (len(track_indices), len(detection_indices))
    )

    for row, track_idx in enumerate(track_indices):
        if tracks[track_idx].time_since_update > 1:
            cost_matrix[row, :] = 1e+5
            continue

        # 计算 IoU 的 cost 值，值越小，IoU 越大
        bbox = tracks[track_idx].to_xywh()
        candidates = np.asarray([detections[i].to_xywh() for i in detection_indices])
        cost_matrix[row, :] = 1. - IoU(bbox, candidates)
    return cost_matrix

